#include <stdio.h>
#include <string.h>
#include <arpa/inet.h>
#include <netinet/ip.h>
#include <stdint.h>
#include <stdlib.h>  // 添加 stdlib.h 头文件
#include "/home/haiyang/wendang/Cprogrammer/include/packet_processing.h"

#define INTERNAL_IP_PREFIX_1 "192.168."
#define INTERNAL_IP_PREFIX_2 "10."
#define INTERNAL_IP_PREFIX_3 "172.16."

// 以太网帧处理函数
void process_ethernet(const uint8_t *packet, int len, TreeNode *next, void *arg) {
    EthernetHeader *eth_header = (EthernetHeader *)packet;
    printf("Source MAC: %02x:%02x:%02x:%02x:%02x:%02x\n",
           eth_header->src_mac[0], eth_header->src_mac[1], eth_header->src_mac[2],
           eth_header->src_mac[3], eth_header->src_mac[4], eth_header->src_mac[5]);
    printf("Destination MAC: %02x:%02x:%02x:%02x:%02x:%02x\n",
           eth_header->dst_mac[0], eth_header->dst_mac[1], eth_header->dst_mac[2],
           eth_header->dst_mac[3], eth_header->dst_mac[4], eth_header->dst_mac[5]);
    if (next != NULL) {
        next->process(packet + sizeof(EthernetHeader), len - sizeof(EthernetHeader), next->left, arg);
    }
}

// IP 头处理函数
void process_ip(const uint8_t *packet, int len, TreeNode *next, void *arg) {
    IPHeader *ip_header = (IPHeader *)packet;
    char src_ip_str[16];
    char dst_ip_str[16];
    inet_ntop(AF_INET, &(ip_header->src_ip), src_ip_str, INET_ADDRSTRLEN);
    inet_ntop(AF_INET, &(ip_header->dst_ip), dst_ip_str, INET_ADDRSTRLEN);
    printf("Source IP: %s\n", src_ip_str);
    printf("Destination IP: %s\n", dst_ip_str);
    TrafficStats *stats = (TrafficStats *)arg;
    PacketFlags flags = {0};

    if (!strstr(src_ip_str, INTERNAL_IP_PREFIX_1) &&
        !strstr(src_ip_str, INTERNAL_IP_PREFIX_2) &&
        !strstr(src_ip_str, INTERNAL_IP_PREFIX_3)) {
        flags.is_external = 1;
    } else if (!strstr(dst_ip_str, INTERNAL_IP_PREFIX_1) &&
               !strstr(dst_ip_str, INTERNAL_IP_PREFIX_2) &&
               !strstr(dst_ip_str, INTERNAL_IP_PREFIX_3)) {
        flags.is_external = 1;
    }

    if (flags.is_external) {
        if (strstr(src_ip_str, INTERNAL_IP_PREFIX_1) ||
            strstr(src_ip_str, INTERNAL_IP_PREFIX_2) ||
            strstr(src_ip_str, INTERNAL_IP_PREFIX_3)) {
            int found = 0;
            for (int i = 0; i < stats->count; i++) {
                if (strcmp(stats->ips[i], src_ip_str) == 0) {
                    stats->traffic[i] += len;
                    found = 1;
                    break;
                }
            }
            if (!found && stats->count < 100) {
                strcpy(stats->ips[stats->count], src_ip_str);
                stats->traffic[stats->count] = len;
                stats->count++;
            }
        } else if (strstr(dst_ip_str, INTERNAL_IP_PREFIX_1) ||
                   strstr(dst_ip_str, INTERNAL_IP_PREFIX_2) ||
                   strstr(dst_ip_str, INTERNAL_IP_PREFIX_3)) {
            int found = 0;
            for (int i = 0; i < stats->count; i++) {
                if (strcmp(stats->ips[i], dst_ip_str) == 0) {
                    stats->traffic[i] += len;
                    found = 1;
                    break;
                }
            }
            if (!found && stats->count < 100) {
                strcpy(stats->ips[stats->count], dst_ip_str);
                stats->traffic[stats->count] = len;
                stats->count++;
            }
        }
    }
}

void init_traffic_stats(TrafficStats *stats) {
    memset(stats->ips, 0, sizeof(stats->ips));
    memset(stats->traffic, 0, sizeof(stats->traffic));
    stats->count = 0;
}

TreeNode *build_tree() {
    TreeNode *ethernet_node = (TreeNode *)malloc(sizeof(TreeNode));
    TreeNode *ip_node = (TreeNode *)malloc(sizeof(TreeNode));

    ethernet_node->process = process_ethernet;
    ethernet_node->left = ip_node;
    ethernet_node->right = NULL;

    ip_node->process = process_ip;
    ip_node->left = NULL;
    ip_node->right = NULL;

    return ethernet_node;
}

void process_packet_through_tree(TreeNode *root, const uint8_t *packet, int len, TrafficStats *stats) {
    if (root != NULL) {
        root->process(packet, len, root->left, stats);
    }
}

void free_tree(TreeNode *root) {
    if (root != NULL) {
        free_tree(root->left);
        free_tree(root->right);
        free(root);
    }
}